package utils

import (
	"context"
	"fmt"
	"github.com/gomodule/redigo/redis"
	"net/url"
	"runtime"
	"sync"
	"time"
)

type RedisPoolMap struct {
	kv map[string]*redis.Pool
	l  sync.RWMutex
}

var rpMap *RedisPoolMap

func init() {
	rpMap = &RedisPoolMap{
		kv: make(map[string]*redis.Pool),
	}
}

// RedisPool 获取一个连接池
// 参考文档: https://pkg.go.dev/github.com/gomodule/redigo/redis
// Redis 命令全集: http://doc.redisfans.com/index.html
func RedisPool(name string, options ...redis.DialOption) (*redis.Pool, error) {
	if name = StrTrim(name); len(name) == 0 {
		return nil, fmt.Errorf("name can not empty")
	}
	dsn := ""
	if u, err := url.Parse(name); err == nil && u.Scheme == "redis" {
		dsn = name
		name = u.Host + ":" + u.Port()
	}
	rpMap.l.RLock()
	rp, ok := rpMap.kv[name]
	rpMap.l.RUnlock()
	if ok && rp != nil {
		return rp, nil
	}
	rpMap.l.Lock()
	defer rpMap.l.Unlock()
	if rp, ok = rpMap.kv[name]; ok && rp != nil {
		return rp, nil
	}
	if dsn == "" {
		if dsn, _ = CfgGetStr(name + "_dsn"); dsn == "" {
			return nil, fmt.Errorf("no config as %s", name+"_dsn")
		}
	}
	cpuNum := runtime.NumCPU()
	maxIdle := cpuNum
	if maxIdle < 20 {
		maxIdle = 20
	}
	maxActive := cpuNum * 5
	if maxActive < 100 {
		maxActive = 100
	}
	rp = &redis.Pool{
		Dial: func() (redis.Conn, error) {
			return redis.DialURL(dsn, options...)
		},
		DialContext: func(ctx context.Context) (redis.Conn, error) {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			default:
				return redis.DialURL(dsn, options...)
			}
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			if time.Since(t) < 3*time.Minute {
				return nil
			}
			_, err := c.Do("PING")
			return err
		},
		MaxIdle:         maxIdle,
		MaxActive:       maxActive,
		IdleTimeout:     3 * time.Minute,
		Wait:            true,
		MaxConnLifetime: 30 * time.Minute,
	}
	if !ok {
		rpMap.kv[name] = rp
	}
	return rp, nil
}

// Copyright (c) 2017 Pavel Pravosud
// https://github.com/rwz/redis-gcra/blob/master/vendor/perform_gcra_ratelimit.lua
var redisRateScript = redis.NewScript(1, `
redis.replicate_commands()
local rate_limit_key = KEYS[1]
local rate = ARGV[1]
local period = ARGV[2]
local emission_interval = period / rate
local burst_offset = emission_interval * rate
local now = redis.call("TIME")
now = (now[1] - 1483228800) + (now[2] / 1000000)
local tat = redis.call("GET", rate_limit_key)
if not tat then
  tat = now
else
  tat = tonumber(tat)
end
tat = math.max(tat, now)
local new_tat = tat + emission_interval
local allow_at = new_tat - burst_offset
local diff = now - allow_at
local remaining = diff / emission_interval
if remaining < 0 then
  return 0
end
local reset_after = new_tat - now
if reset_after > 0 then
  redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
end
return 1
`)

// RedisRateLimit Redis的令牌桶方式限流
func RedisRateLimit(rc redis.Conn, key string, rate uint, period time.Duration) (bool, error) {
	return redis.Bool(redisRateScript.Do(rc, key, rate, period.Seconds()))
}

// RedisSet 写数据，支持 px/nx/xx
func RedisSet(rc redis.Conn, key string, value string, period time.Duration, NXorXX string) (bool, error) {
	if bln, err := InArray(NXorXX, []string{"", "NX", "XX"}); !bln || err != nil {
		return false, fmt.Errorf("[NXorXX] should be \"\" or \"NX\" or \"XX\"")
	}
	var args []interface{}
	args = append(args, key, value)
	if period > 0 {
		args = append(args, "PX", period.Milliseconds())
	}
	if NXorXX != "" {
		args = append(args, NXorXX)
	}
	ok, err := redis.String(rc.Do("SET", args...))
	return ok == "OK", err
}

// RedisIncr Redis 计数器
func RedisIncr(rc redis.Conn, key string, step int) (int64, error) {
	return redis.Int64(rc.Do("INCRBY", key, step))
}

// RedisDel Redis 删除一个或多个 Key
func RedisDel(rc redis.Conn, key ...interface{}) (int64, error) {
	return redis.Int64(rc.Do("DEL", key...))
}

// RedisExpire Redis 设置过期时间
func RedisExpire(rc redis.Conn, key string, period time.Duration) (bool, error) {
	return redis.Bool(rc.Do("PEXPIRE", key, period.Milliseconds()))
}
